from xtquant import xtdata
import pandas as pd
from tqdm import tqdm
from loguru import logger
import feather
import os
import time


logger.add('runtime_{time}.log', rotation="3 MB", retention='15 days')


def get_trade_cal(market='SH'):
    # 读取交易日历
    trade_cal = xtdata.get_trading_dates(market=market, start_time='', end_time='', count=-1)
    cal_ = pd.to_datetime(trade_cal, unit='ms') + pd.to_timedelta('8H')
    return cal_

def get_market_data(symbol, period_, count_=-1, start_time='20190101'):
    field_list = ['time', 'open', 'high', 'low', 'close', 'volume', 'amount', 'preClose']

    hist = xtdata.get_market_data(field_list, stock_list=[symbol],
                                    period=period_, start_time=start_time, end_time='', count=count_,
                                    dividend_type='none', fill_data=True)

    df = pd.DataFrame({key: value.values[0] for key, value in hist.items()})
    df['time'] = pd.to_datetime(df['time'], unit='ms') + pd.to_timedelta('8H')

    if len(df) == 0:
        return pd.DataFrame()
    df.loc[:, 'symbol'] = symbol
    df = df.set_index('time')    # set index to time, then drop the original column 'time'
    return df

def get_full_data(symbol, res_path, start_dt, period_='1d'):
    df_dt = get_market_data(symbol=symbol, period_=period_, start_time=start_dt)
    if len(df_dt) > 0:
        feather.write_dataframe(df_dt, os.path.join(res_path, f'{symbol}_{period_}.feather'))
        print(f"{symbol} 保存完成")
        return df_dt
    else:
        print(f"{symbol} 数据未获取到，不保存")
        return pd.DataFrame()
    
def save_1d_dt(symbol, start_dt, periods_='1d'):
    xtdata.download_history_data(stock_code=symbol, period=periods_, start_time=start_dt, end_time='')

def update_etf_stk_cbond(start_dt, periods_='1d'):
    # 存储ETF
    etfhs = xtdata.get_stock_list_in_sector(sector_name='沪深ETF')

    t1 = time.time()
    for etf in tqdm(etfhs, desc=f'prepare_etf_{periods_}'):
        save_1d_dt(symbol=etf, start_dt=start_dt, periods_=periods_)
    logger.info(f'prepare_etf_{periods_} done {t1-time.time()}')

    # 存储股票
    stks = xtdata.get_stock_list_in_sector(sector_name='沪深A股')

    t1 = time.time()
    for etf in tqdm(stks, desc=f'prepare_stk_{periods_}'):
        save_1d_dt(symbol=etf, start_dt=start_dt, periods_=periods_)
    logger.info(f'prepare_stk_{periods_} done {t1-time.time()}')

    # 存储可转债
    cbond = xtdata.get_stock_list_in_sector(sector_name='沪深转债')

    t1 = time.time()
    for exbond in tqdm(cbond, desc=f'prepare_cbond_{periods_}'):
        save_1d_dt(symbol=exbond, start_dt=start_dt, periods_=periods_)
    logger.info(f'prepare_cbond_{periods_} done {t1-time.time()}')


if __name__ == '__main__':
    start_t = '20200101'
    update_etf_stk_cbond(start_dt=start_t)
